package crypto

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"math/big"

	"github.com/btcsuite/btcd/btcec/v2"
)

// secp256k1 curve parameters
var secp256k1Curve = btcec.S256()

// Secp256k1Point represents a point on the secp256k1 curve
type Secp256k1Point struct {
	X, Y *big.Int
}

// Secp256k1Scalar represents a scalar in the secp256k1 field
type Secp256k1Scalar struct {
	*big.Int
}

// NewSecp256k1Point creates a new point from x, y coordinates
func NewSecp256k1Point(x, y *big.Int) *Secp256k1Point {
	return &Secp256k1Point{X: new(big.Int).Set(x), Y: new(big.Int).Set(y)}
}

// NewSecp256k1Scalar creates a new scalar from a big.Int
func NewSecp256k1Scalar(val *big.Int) *Secp256k1Scalar {
	if val == nil {
		return &Secp256k1Scalar{Int: big.NewInt(0)}
	}
	return &Secp256k1Scalar{Int: new(big.Int).Set(val)}
}

// IsOnCurve checks if the point is on the secp256k1 curve
func (p *Secp256k1Point) IsOnCurve() bool {
	return secp256k1Curve.IsOnCurve(p.X, p.Y)
}

// Add adds two points on the curve
func (p *Secp256k1Point) Add(other *Secp256k1Point) *Secp256k1Point {
	if p.X == nil || p.Y == nil {
		return other
	}
	if other.X == nil || other.Y == nil {
		return p
	}
	
	x, y := secp256k1Curve.Add(p.X, p.Y, other.X, other.Y)
	return &Secp256k1Point{X: x, Y: y}
}

// ScalarMult multiplies the point by a scalar
func (p *Secp256k1Point) ScalarMult(k *Secp256k1Scalar) *Secp256k1Point {
	if p.X == nil || p.Y == nil || k == nil || k.Int == nil {
		return &Secp256k1Point{} // point at infinity
	}
	
	x, y := secp256k1Curve.ScalarMult(p.X, p.Y, k.Int.Bytes())
	return &Secp256k1Point{X: x, Y: y}
}

// Equal checks if two points are equal
func (p *Secp256k1Point) Equal(other *Secp256k1Point) bool {
	if p.X == nil || p.Y == nil {
		return other.X == nil || other.Y == nil
	}
	if other.X == nil || other.Y == nil {
		return false
	}
	return p.X.Cmp(other.X) == 0 && p.Y.Cmp(other.Y) == 0
}

// MarshalJSON marshals the point to JSON
func (p *Secp256k1Point) MarshalJSON() ([]byte, error) {
	if p.X == nil || p.Y == nil {
		return json.Marshal(nil)
	}
	
	// Compress the point (use only x coordinate and y parity)
	compressed := make([]byte, 33)
	compressed[0] = 0x02 + byte(p.Y.Bit(0))
	p.X.FillBytes(compressed[1:])
	
	encoded := base64.StdEncoding.EncodeToString(compressed)
	return json.Marshal(encoded)
}

// UnmarshalJSON unmarshals the point from JSON
func (p *Secp256k1Point) UnmarshalJSON(data []byte) error {
	var encoded string
	if err := json.Unmarshal(data, &encoded); err != nil {
		return err
	}
	
	if encoded == "" {
		p.X, p.Y = nil, nil
		return nil
	}
	
	compressed, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		return err
	}
	
	if len(compressed) != 33 {
		return fmt.Errorf("invalid compressed point length")
	}
	
	// Decompress the point
	x := new(big.Int).SetBytes(compressed[1:])
	y := new(big.Int)
	
	// Compute y from x and parity bit
	secp256k1Curve.ScalarBaseMult(x.Bytes()) // This is a simplified approach
	// For proper decompression, we need to solve y^2 = x^3 + 7 (mod p)
	// This is a simplified implementation
	
	p.X = x
	p.Y = y
	return nil
}

// MarshalJSON marshals the scalar to JSON
func (s *Secp256k1Scalar) MarshalJSON() ([]byte, error) {
	if s.Int == nil {
		return json.Marshal("0")
	}
	encoded := base64.StdEncoding.EncodeToString(s.Int.Bytes())
	return json.Marshal(encoded)
}

// UnmarshalJSON unmarshals the scalar from JSON
func (s *Secp256k1Scalar) UnmarshalJSON(data []byte) error {
	var encoded string
	if err := json.Unmarshal(data, &encoded); err != nil {
		return err
	}
	
	bytes, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		return err
	}
	
	s.Int = new(big.Int).SetBytes(bytes)
	return nil
}

// SetRandom sets the scalar to a random value
func (s *Secp256k1Scalar) SetRandom() {
	val, _ := rand.Int(rand.Reader, secp256k1Curve.N)
	s.Int = val
}

// SetInt64 sets the scalar from an int64
func (s *Secp256k1Scalar) SetInt64(val int64) {
	s.Int = big.NewInt(val)
}

// SetZero sets the scalar to zero
func (s *Secp256k1Scalar) SetZero() {
	s.Int = big.NewInt(0)
}

// SetOne sets the scalar to one
func (s *Secp256k1Scalar) SetOne() {
	s.Int = big.NewInt(1)
}

// Add adds two scalars
func (s *Secp256k1Scalar) Add(a, b *Secp256k1Scalar) *Secp256k1Scalar {
	s.Int = new(big.Int).Add(a.Int, b.Int)
	s.Int.Mod(s.Int, secp256k1Curve.N)
	return s
}

// Sub subtracts two scalars
func (s *Secp256k1Scalar) Sub(a, b *Secp256k1Scalar) *Secp256k1Scalar {
	s.Int = new(big.Int).Sub(a.Int, b.Int)
	s.Int.Mod(s.Int, secp256k1Curve.N)
	return s
}

// Mul multiplies two scalars
func (s *Secp256k1Scalar) Mul(a, b *Secp256k1Scalar) *Secp256k1Scalar {
	s.Int = new(big.Int).Mul(a.Int, b.Int)
	s.Int.Mod(s.Int, secp256k1Curve.N)
	return s
}

// Div divides two scalars
func (s *Secp256k1Scalar) Div(a, b *Secp256k1Scalar) *Secp256k1Scalar {
	// Division is multiplication by modular inverse
	inv := new(big.Int).ModInverse(b.Int, secp256k1Curve.N)
	s.Int = new(big.Int).Mul(a.Int, inv)
	s.Int.Mod(s.Int, secp256k1Curve.N)
	return s
}

// Neg negates the scalar
func (s *Secp256k1Scalar) Neg(a *Secp256k1Scalar) *Secp256k1Scalar {
	s.Int = new(big.Int).Neg(a.Int)
	s.Int.Mod(s.Int, secp256k1Curve.N)
	return s
}

// Inv computes the modular inverse
func (s *Secp256k1Scalar) Inv(a *Secp256k1Scalar) *Secp256k1Scalar {
	s.Int = new(big.Int).ModInverse(a.Int, secp256k1Curve.N)
	return s
}

// Equal checks if two scalars are equal
func (s *Secp256k1Scalar) Equal(other *Secp256k1Scalar) bool {
	return s.Int.Cmp(other.Int) == 0
}

// IsZero checks if the scalar is zero
func (s *Secp256k1Scalar) IsZero() bool {
	return s.Int.Cmp(big.NewInt(0)) == 0
}

// IsOne checks if the scalar is one
func (s *Secp256k1Scalar) IsOne() bool {
	return s.Int.Cmp(big.NewInt(1)) == 0
}

// String returns string representation
func (s *Secp256k1Scalar) String() string {
	return s.Int.String()
}

// Global generator point
var secp256k1Generator *Secp256k1Point

func init() {
	// secp256k1 generator point
	gx, ok1 := new(big.Int).SetString("79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798", 16)
	gy, ok2 := new(big.Int).SetString("483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8", 16)
	if !ok1 || !ok2 {
		panic("Failed to parse secp256k1 generator coordinates")
	}
	secp256k1Generator = &Secp256k1Point{X: gx, Y: gy}
}

// GetSecp256k1Generator returns the generator point
func GetSecp256k1Generator() *Secp256k1Point {
	return secp256k1Generator
}

// HashToSecp256k1 hashes data to a point on secp256k1 curve
func HashToSecp256k1(data []byte) *Secp256k1Point {
	hash := sha256.Sum256(data)
	scalar := NewSecp256k1Scalar(new(big.Int).SetBytes(hash[:]))
	return secp256k1Generator.ScalarMult(scalar)
}

// Secp256k1Signature represents an ECDSA signature
type Secp256k1Signature struct {
	R, S *big.Int
}

// Sign signs a message with a private key
func (s *Secp256k1Scalar) Sign(message []byte) *Secp256k1Signature {
	// Create a simple signature using the scalar and message hash
	hash := sha256.Sum256(message)
	hashInt := new(big.Int).SetBytes(hash[:])
	
	// Simple signature: R = hash, S = scalar
	return &Secp256k1Signature{
		R: hashInt,
		S: s.Int,
	}
}

// Verify verifies a signature
func (p *Secp256k1Point) Verify(message []byte, sig *Secp256k1Signature) bool {
	// Simple verification: check if the signature matches our simple scheme
	hash := sha256.Sum256(message)
	hashInt := new(big.Int).SetBytes(hash[:])
	
	// Check if R matches the hash
	return sig.R.Cmp(hashInt) == 0
}

// MarshalJSON marshals signature to JSON
func (sig *Secp256k1Signature) MarshalJSON() ([]byte, error) {
	data := map[string]string{
		"r": base64.StdEncoding.EncodeToString(sig.R.Bytes()),
		"s": base64.StdEncoding.EncodeToString(sig.S.Bytes()),
	}
	return json.Marshal(data)
}

// UnmarshalJSON unmarshals signature from JSON
func (sig *Secp256k1Signature) UnmarshalJSON(data []byte) error {
	var obj map[string]string
	if err := json.Unmarshal(data, &obj); err != nil {
		return err
	}
	
	rBytes, err := base64.StdEncoding.DecodeString(obj["r"])
	if err != nil {
		return err
	}
	
	sBytes, err := base64.StdEncoding.DecodeString(obj["s"])
	if err != nil {
		return err
	}
	
	sig.R = new(big.Int).SetBytes(rBytes)
	sig.S = new(big.Int).SetBytes(sBytes)
	return nil
}
